"""Baidu Wenxin (文心一言) model adapter implementation."""

from typing import AsyncGenerator, Dict, Any, List
import aiohttp
import time
import hashlib
import json

from .base import BaseModelAdapter, ModelConfig, ModelResponse


class WenxinAdapter(BaseModelAdapter):
    """Adapter for Baidu Wenxin models."""
    
    # Pricing per 1K tokens
    PRICING = {
        "ernie-bot-4": {"prompt": 0.012, "completion": 0.012},
        "ernie-bot": {"prompt": 0.004, "completion": 0.004},
        "ernie-bot-turbo": {"prompt": 0.002, "completion": 0.002}
    }
    
    def __init__(self, config: ModelConfig):
        """Initialize with access token management."""
        super().__init__(config)
        self._access_token = None
        self._token_expiry = 0
    
    def _initialize_client(self) -> None:
        """Initialize Wenxin client."""
        # Wenxin uses access token, not direct API key
        self._api_key = self.config.api_key  # This should be app_key
        self._secret_key = self.config.extra_params.get("secret_key", "")
        self._base_url = self.config.api_base or "https://aip.baidubce.com"
    
    async def _get_access_token(self) -> str:
        """Get or refresh access token.
        
        Returns:
            Valid access token
        """
        if self._access_token and time.time() < self._token_expiry:
            return self._access_token
        
        # Refresh token
        url = f"{self._base_url}/oauth/2.0/token"
        params = {
            "grant_type": "client_credentials",
            "client_id": self._api_key,
            "client_secret": self._secret_key
        }
        
        async with aiohttp.ClientSession() as session:
            async with session.get(url, params=params) as response:
                data = await response.json()
                if "access_token" in data:
                    self._access_token = data["access_token"]
                    # Token expires in 30 days, refresh after 29 days
                    self._token_expiry = time.time() + (29 * 24 * 3600)
                    return self._access_token
                else:
                    raise Exception(f"Failed to get access token: {data}")
    
    async def generate(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> ModelResponse:
        """Generate response from Wenxin model.
        
        Args:
            messages: Conversation messages
            **kwargs: Additional parameters
            
        Returns:
            ModelResponse with generated content
        """
        access_token = await self._get_access_token()
        
        # Wenxin API endpoint
        model_endpoints = {
            "ernie-bot-4": "completions_pro",
            "ernie-bot": "completions",
            "ernie-bot-turbo": "eb-instant"
        }
        
        endpoint = model_endpoints.get(self.config.model_name, "completions")
        url = f"{self._base_url}/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/{endpoint}"
        
        headers = {
            "Content-Type": "application/json"
        }
        
        params = {
            "access_token": access_token
        }
        
        data = {
            "messages": messages,
            "temperature": self.config.temperature,
            "top_p": self.config.top_p,
            "stream": False
        }
        
        data.update(kwargs)
        data.update(self.config.extra_params)
        
        async with aiohttp.ClientSession() as session:
            async with session.post(
                url,
                headers=headers,
                params=params,
                json=data,
                timeout=aiohttp.ClientTimeout(total=self.config.timeout)
            ) as response:
                result = await response.json()
                
                if "error_code" in result:
                    raise Exception(f"Wenxin API error: {result.get('error_msg', 'Unknown error')}")
                
                usage = {
                    "prompt_tokens": result.get("usage", {}).get("prompt_tokens", 0),
                    "completion_tokens": result.get("usage", {}).get("completion_tokens", 0),
                    "total_tokens": result.get("usage", {}).get("total_tokens", 0)
                }
                
                return ModelResponse(
                    content=result.get("result", ""),
                    model=self.config.model_name,
                    usage=usage,
                    metadata={
                        "id": result.get("id", ""),
                        "is_truncated": result.get("is_truncated", False)
                    }
                )
    
    async def stream(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> AsyncGenerator[str, None]:
        """Stream response from Wenxin model."""
        # Similar to generate but with stream=True
        # Implementation would follow SSE pattern
        yield "Wenxin streaming not fully implemented yet"
    
    def count_tokens(self, text: str) -> int:
        """Count tokens for Wenxin models."""
        # Rough estimate for Chinese text
        return len(text)
    
    def estimate_cost(self, prompt_tokens: int, completion_tokens: int) -> float:
        """Estimate cost for Wenxin models."""
        model_key = self.config.model_name
        
        pricing = self.PRICING.get(model_key, self.PRICING["ernie-bot-turbo"])
        prompt_cost = (prompt_tokens / 1000) * pricing["prompt"]
        completion_cost = (completion_tokens / 1000) * pricing["completion"]
        
        # Convert CNY to USD
        cny_to_usd = 0.14
        return (prompt_cost + completion_cost) * cny_to_usd